package com.wt.admin.config;

import com.wt.admin.config.cache.Cache;
import com.wt.admin.config.cache.CacheManager;
import com.wt.admin.config.cache.impl.ChatContentCache;
import com.wt.admin.config.prop.AIProp;
import com.wt.admin.domain.vo.ai.ChatModelContentVO;
import com.wt.admin.service.vector.Vector;
import com.wt.admin.service.vector.impl.ESVectorImpl;
import com.wt.admin.service.vector.impl.MemoryVectorImpl;
import io.micrometer.observation.ObservationRegistry;
import jakarta.annotation.Resource;
import lombok.Data;
import org.dromara.easyai.config.SentenceConfig;
import org.dromara.easyai.config.TfConfig;
import org.dromara.easyai.entity.KeyWordForSentence;
import org.dromara.easyai.naturalLanguage.TalkToTalk;
import org.dromara.easyai.naturalLanguage.languageCreator.CatchKeyWord;
import org.dromara.easyai.naturalLanguage.word.MyKeyWord;
import org.dromara.easyai.naturalLanguage.word.WordEmbedding;
import org.dromara.easyai.rnnJumpNerveCenter.RRNerveManager;
import org.dromara.easyai.yolo.FastYolo;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.ollama.management.ModelManagementOptions;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.openai.OpenAiEmbeddingOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.web.servlet.FilterRegistrationBean;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.filter.RequestContextFilter;

import java.util.ArrayList;
import java.util.List;

@Configuration
public class AIConfig {

    @Resource
    private AIProp aiProp;

    /**
     * 词向量嵌入器,语义神经
     * @return
     */
    @Bean("wordEmbedding")
    public Cache<String,WordAndRRManager> wordEmbedding(){
        return CacheManager.getCache("wordEmbedding");
    }

    @Bean("myKeyWord")
    public Cache<Integer, MyKeyWord> myKeyWordMap(){
        return CacheManager.getCache("myKeyWord");
    }

    @Bean("catchKeyWord")
    public Cache<Integer, CatchKeyWord> catchKeyWordMap(){
        return CacheManager.getCache("catchKeyWord");
    }

    @Bean("sensorKeyWordMapper")
    public Cache<Integer,Cache<Integer, List<KeyWordForSentence>>> sensorKeyWordMapper(){
        return CacheManager.getCache("sensorKeyWordMapper");
    }

    @Bean("imageYoloManager>")
    public Cache<String, FastYolo> imageYoloManager(){
        return CacheManager.getCache("imageYoloManager");
    }

    @Bean("keyWordValueList")
    public List<KeywordValue> keyWordValueList(){
        return new ArrayList<>();
    }

    /**
     * 缓存聊天内容
     * @return
     */
    @Bean("chatContents")
    public ChatContentCache<Integer, ChatModelContentVO> chatContents(){
        return new ChatContentCache<>();
    }

    @Bean("embeddingModel")
    @ConditionalOnProperty(name = "spring.embedding.method", havingValue = "ollama")
    public EmbeddingModel ollamaEmbeddingModel(){
        OllamaApi build = OllamaApi.builder().baseUrl(aiProp.getEmbedding().getUrl()).build();
        return new OllamaEmbeddingModel(
                build,
                OllamaOptions.builder()
                        .model(aiProp.getEmbedding().getModel())
                        .build(),
                ObservationRegistry.NOOP,
                ModelManagementOptions.defaults()
        );
    }

    @Bean("embeddingModel")
    @ConditionalOnProperty(name = "spring.embedding.method", havingValue = "openai")
    public EmbeddingModel openAiEmbeddingModel(){
        var openAiApi = OpenAiApi.builder()
                .baseUrl(aiProp.getEmbedding().getUrl())
                .apiKey(aiProp.getEmbedding().getApiKey())
                .build();
        return new OpenAiEmbeddingModel(
                openAiApi,
                MetadataMode.EMBED,
                OpenAiEmbeddingOptions.builder()
                        .model(aiProp.getEmbedding().getModel())
                        .build(),
                RetryUtils.DEFAULT_RETRY_TEMPLATE);
    }

    @Bean
    public SimpleVectorStore simpleVectorStore(@Qualifier("embeddingModel") EmbeddingModel embeddingModel) {
        return SimpleVectorStore.builder(embeddingModel)
                .batchingStrategy(new TokenCountBatchingStrategy())
                .build();
    }

//    @Bean
//    @ConditionalOnProperty(name = "spring.vector.method", havingValue = "es")
//    public ElasticsearchVectorStore elasticsearchVectorStore(EmbeddingModel embeddingModel, ElasticsearchProperties elasticsearchProperties) {
//        HttpHost[] hosts = elasticsearchProperties.getUris().stream()
//                .map(HttpHost::create)
//                .toArray(HttpHost[]::new);
//        RestClientBuilder builder = RestClient.builder(hosts);
//
//        builder.setHttpClientConfigCallback(httpClientBuilder -> {
//            try {
//                // 创建信任所有证书的SSL上下文 禁用SSL验证
//                SSLContext sslContext = SSLContextBuilder
//                        .create()
//                        .loadTrustMaterial(null, (chain, authType) -> true)
//                        .build();
//
//                return httpClientBuilder
//                        .setSSLContext(sslContext)
//                        .setSSLHostnameVerifier((hostname, session) -> true)
//                        .setKeepAliveStrategy((response, context) -> {
//                            // 设置keep-alive时间为30秒
//                            return 30000;
//                        })
//                        .setMaxConnTotal(20)// 配置连接池和超时
//                        .setMaxConnPerRoute(20)
//                        .setConnectionTimeToLive(60, java.util.concurrent.TimeUnit.SECONDS)
//                        .setDefaultRequestConfig(
//                                org.apache.http.client.config.RequestConfig.custom()
//                                        .setConnectTimeout(10000)  // 10秒连接超时
//                                        .setSocketTimeout(30000)   // 30秒Socket超时
//                                        .build()
//                        );
//            } catch (Exception e) {
//                throw new RuntimeException("Failed to create SSL context", e);
//            }
//        });
//        // 使用Basic认证方式
//        String auth = elasticsearchProperties.getUsername() + ":" + elasticsearchProperties.getPassword();
//        String encodedAuth = java.util.Base64.getEncoder().encodeToString(auth.getBytes());
//        builder.setDefaultHeaders(new Header[]{
//                new BasicHeader("Authorization", "Basic " + encodedAuth),
//                new BasicHeader("Connection", "keep-alive")
//        });
//
//        ElasticsearchVectorStoreOptions options = new ElasticsearchVectorStoreOptions();
//        options.setIndexName("custom-index");    // Optional: defaults to "spring-ai-document-index"
//        options.setSimilarity(SimilarityFunction.cosine);           // Optional: defaults to COSINE
//        options.setDimensions(1536);             // Optional: defaults to model dimensions or 1536
//
//        return ElasticsearchVectorStore.builder(builder.build(), embeddingModel)
//                .options(options)                     // Optional: use custom options
//                .initializeSchema(true)               // Optional: defaults to false
//                .batchingStrategy(new TokenCountBatchingStrategy()) // Optional: defaults to TokenCountBatchingStrategy
//                .build();
//    }

    @Bean
    @ConditionalOnProperty(name = "spring.vector.method", havingValue = "es")
    public Vector es() {
        return new ESVectorImpl();
    }

    @Bean
    @ConditionalOnProperty(name = "spring.vector.method", havingValue = "memory")
    public Vector memory() {
        return new MemoryVectorImpl();
    }

    @Bean
    public ChatMemory chatMemory() {
        return  MessageWindowChatMemory.builder()
                .maxMessages(10)
                .chatMemoryRepository(new InMemoryChatMemoryRepository())
                .build();
    }


    @Bean
    public FilterRegistrationBean<RequestContextFilter> requestContextFilter() {
        FilterRegistrationBean<RequestContextFilter> registrationBean = new FilterRegistrationBean<>();
        registrationBean.setFilter(new RequestContextFilter());
        registrationBean.setOrder(1);
        registrationBean.addUrlPatterns("/*");
        return registrationBean;
    }


    @Data
    public static class KeywordValue{
        private String keywordValue;
        private Integer typeId;
        private Long index;//该关键词索引id
        public KeywordValue(String keywordValue, Integer typeId, Long index) {
            this.keywordValue = keywordValue;
            this.typeId = typeId;
            this.index = index;
        }
        public KeywordValue(String keywordValue, Integer typeId) {
            this.keywordValue = keywordValue;
            this.typeId = typeId;
        }
    }

    @Data
    public static class WordAndRRManager {

        private WordEmbedding wordEmbedding;
        private RRNerveManager rrNerveManager;
        private TalkToTalk talkToTalk;
        private SentenceConfig sentenceConfig;
        private TfConfig tfConfig;


        public WordAndRRManager() {
            this.wordEmbedding = new WordEmbedding();
            this.rrNerveManager = new RRNerveManager(wordEmbedding);
        }

    }
}
